"""
Comprehensive tests for VectorField and Tensor2Field implementations

Tests the complete vector and tensor field system including:
- Construction and basic operations
- Component access and manipulation  
- Vector operations (dot, cross, magnitude)
- Tensor operations (trace, determinant, eigenvalues)
- Vector calculus (divergence, curl, gradient)
- Field collections
- Mathematical correctness
"""

using Test
using LinearAlgebra
using Statistics
using PDEJulia
using PDEJulia.Grids
using PDEJulia.Fields

@testset "Vector and Tensor Field System Tests" begin

@testset "VectorField Construction and Basic Operations" begin
    # Create test grids
    grid_1d = UnitGrid([21])
    grid_2d = UnitGrid([21, 21]) 
    grid_3d = UnitGrid([11, 11, 11])
    
    @testset "Construction" begin
        # Test construction from function
        field_2d = VectorField(pos -> [sin(pos[1]), cos(pos[2])], grid_2d)
        @test field_2d isa VectorField{Float64,2}
        @test grid(field_2d) === grid_2d
        @test size(data(field_2d)) == (21, 21, 2)
        
        # Test zero construction
        zero_field = VectorField(Float64, grid_2d)
        @test all(data(zero_field) .== 0)
        
        # Test construction from scalar fields
        fx = ScalarField(pos -> sin(pos[1]), grid_2d)
        fy = ScalarField(pos -> cos(pos[2]), grid_2d)
        from_scalars = VectorField([fx, fy])
        @test from_scalars isa VectorField{Float64,2}
        
        # Test dimension consistency
        @test_throws DimensionError VectorField([fx])  # Wrong number of components
    end
    
    @testset "Component Access" begin
        field_2d = VectorField(pos -> [pos[1]^2, pos[2]^3], grid_2d)
        
        # Test integer indexing
        x_comp = field_2d[1]
        y_comp = field_2d[2]
        @test x_comp isa ScalarField
        @test y_comp isa ScalarField
        @test_throws BoundsError field_2d[3]
        
        # Test component modification
        field_2d[1] = 2.0
        @test all(data(field_2d)[:, :, 1] .== 2.0)
        
        # Test component replacement with scalar field
        new_comp = ScalarField(pos -> exp(pos[1]), grid_2d)
        field_2d[2] = new_comp
        @test data(field_2d)[:, :, 2] ≈ data(new_comp)
    end
    
    @testset "Basic Interface" begin
        field_2d = VectorField(Float64, grid_2d)
        
        # Test similar and copy
        sim_field = similar(field_2d)
        @test typeof(sim_field) == typeof(field_2d)
        @test grid(sim_field) === grid(field_2d)
        
        copied_field = copy(field_2d)
        @test data(copied_field) == data(field_2d)
        @test data(copied_field) !== data(field_2d)  # Different memory
        
        # Test size and length
        @test size(field_2d) == (21, 21)
        @test ndims(field_2d) == 2
        @test eltype(field_2d) == Float64
    end
end

@testset "Vector Operations" begin
    grid_2d = UnitGrid([21, 21])
    grid_3d = UnitGrid([11, 11, 11])
    
    @testset "Dot Product" begin
        v1 = VectorField(pos -> [1.0, 2.0], grid_2d)
        v2 = VectorField(pos -> [3.0, 4.0], grid_2d) 
        
        dot_result = dot(v1, v2)
        @test dot_result isa ScalarField
        @test all(data(dot_result) .≈ 11.0)  # 1*3 + 2*4 = 11
        
        # Test incompatible grids
        v3 = VectorField(Float64, UnitGrid([11, 11]))
        @test_throws ArgumentError dot(v1, v3)
    end
    
    @testset "Cross Product" begin
        v1 = VectorField(pos -> [1.0, 0.0, 0.0], grid_3d)
        v2 = VectorField(pos -> [0.0, 1.0, 0.0], grid_3d)
        
        cross_result = cross(v1, v2)
        @test cross_result isa VectorField{Float64,3}
        
        # Should give [0, 0, 1] everywhere
        @test all(data(cross_result)[:, :, :, 1] .≈ 0.0)
        @test all(data(cross_result)[:, :, :, 2] .≈ 0.0)
        @test all(data(cross_result)[:, :, :, 3] .≈ 1.0)
    end
    
    @testset "Magnitude and Normalization" begin
        v = VectorField(pos -> [3.0, 4.0], grid_2d)
        
        mag = magnitude(v)
        @test mag isa ScalarField
        @test all(data(mag) .≈ 5.0)  # sqrt(3^2 + 4^2) = 5
        
        # Test norm alias
        @test norm(v) == mag
        
        # Test magnitude squared
        mag_sq = magnitude_squared(v)
        @test all(data(mag_sq) .≈ 25.0)
        
        # Test normalization
        v_norm = normalize(v)
        norm_mag = magnitude(v_norm)
        @test all(data(norm_mag) .≈ 1.0)
        
        # Test in-place normalization
        v_copy = copy(v)
        normalize!(v_copy)
        @test all(data(magnitude(v_copy)) .≈ 1.0)
    end
end

@testset "Vector Calculus Operations" begin
    grid_2d = UnitGrid([21, 21])
    
    @testset "Divergence" begin
        # Test with simple linear field: v = [x, y]
        # Divergence should be 2 everywhere
        v = VectorField(pos -> [pos[1], pos[2]], grid_2d)
        div_v = divergence(v)
        
        @test div_v isa ScalarField
        # Note: Actual test depends on grid_gradient implementation
        # This tests the interface and basic functionality
    end
    
    @testset "2D Curl" begin
        # Test 2D curl (returns scalar)
        v = VectorField(pos -> [pos[2], -pos[1]], grid_2d)
        curl_v = curl(v)
        
        @test curl_v isa ScalarField
        # For v = [y, -x], curl should be -2 everywhere
    end
    
    @testset "Gradient" begin
        # Gradient of vector field returns tensor field
        v = VectorField(pos -> [pos[1]^2, pos[2]^2], grid_2d)
        grad_v = gradient(v)
        
        @test grad_v isa Tensor2Field
        @test size(data(grad_v)) == (21, 21, 2, 2)
    end
    
    @testset "Laplacian" begin
        v = VectorField(pos -> [sin(pos[1]), cos(pos[2])], grid_2d)
        lap_v = laplacian(v)
        
        @test lap_v isa VectorField{Float64,2}
        @test size(data(lap_v)) == size(data(v))
    end
end

@testset "Vector Field Mathematical Operations" begin
    grid_2d = UnitGrid([21, 21])
    
    v1 = VectorField(pos -> [1.0, 2.0], grid_2d)
    v2 = VectorField(pos -> [3.0, 4.0], grid_2d)
    scalar_field = ScalarField(pos -> 2.0, grid_2d)
    
    @testset "Element-wise Operations" begin
        # Addition
        v_add = v1 + v2
        @test all(data(v_add)[:, :, 1] .≈ 4.0)
        @test all(data(v_add)[:, :, 2] .≈ 6.0)
        
        # Subtraction
        v_sub = v2 - v1
        @test all(data(v_sub)[:, :, 1] .≈ 2.0)
        @test all(data(v_sub)[:, :, 2] .≈ 2.0)
        
        # Scalar multiplication
        v_mult = v1 * 3.0
        @test all(data(v_mult)[:, :, 1] .≈ 3.0)
        @test all(data(v_mult)[:, :, 2] .≈ 6.0)
        
        # Scalar division
        v_div = v2 / 2.0
        @test all(data(v_div)[:, :, 1] .≈ 1.5)
        @test all(data(v_div)[:, :, 2] .≈ 2.0)
        
        # Unary negation
        v_neg = -v1
        @test all(data(v_neg)[:, :, 1] .≈ -1.0)
        @test all(data(v_neg)[:, :, 2] .≈ -2.0)
    end
    
    @testset "Scalar Field Multiplication" begin
        v_scaled = scalar_field * v1
        @test all(data(v_scaled)[:, :, 1] .≈ 2.0)
        @test all(data(v_scaled)[:, :, 2] .≈ 4.0)
        
        # Test commutativity
        v_scaled2 = v1 * scalar_field
        @test data(v_scaled) ≈ data(v_scaled2)
    end
    
    @testset "In-place Operations" begin
        v_copy = copy(v1)
        v_copy += v2
        @test all(data(v_copy)[:, :, 1] .≈ 4.0)
        
        v_copy = copy(v1)
        v_copy *= 2.0
        @test all(data(v_copy)[:, :, 1] .≈ 2.0)
    end
end

@testset "Tensor2Field Construction and Basic Operations" begin
    grid_2d = UnitGrid([21, 21])
    
    @testset "Construction" begin
        # Test construction from function
        tensor_field = Tensor2Field(pos -> [pos[1] pos[2]; pos[2] pos[1]], grid_2d)
        @test tensor_field isa Tensor2Field{Float64,2}
        @test size(data(tensor_field)) == (21, 21, 2, 2)
        
        # Test zero construction
        zero_tensor = Tensor2Field(Float64, grid_2d)
        @test all(data(zero_tensor) .== 0)
        
        # Test identity tensor
        identity_tensor = identity_tensor(grid_2d)
        @test all(data(identity_tensor)[:, :, 1, 1] .== 1.0)
        @test all(data(identity_tensor)[:, :, 1, 2] .== 0.0)
        @test all(data(identity_tensor)[:, :, 2, 1] .== 0.0)
        @test all(data(identity_tensor)[:, :, 2, 2] .== 1.0)
    end
    
    @testset "Component Access" begin
        tensor_field = Tensor2Field(pos -> [1.0 2.0; 3.0 4.0], grid_2d)
        
        # Test component access
        t11 = tensor_field[1,1]
        t12 = tensor_field[1,2]
        t21 = tensor_field[2,1]
        t22 = tensor_field[2,2]
        
        @test t11 isa ScalarField
        @test all(data(t11) .== 1.0)
        @test all(data(t12) .== 2.0)
        @test all(data(t21) .== 3.0)
        @test all(data(t22) .== 4.0)
        
        # Test bounds checking
        @test_throws BoundsError tensor_field[3,1]
        @test_throws BoundsError tensor_field[1,3]
    end
    
    @testset "Matrix Operations" begin
        tensor_field = Tensor2Field(pos -> [1.0 2.0; 3.0 4.0], grid_2d)
        
        # Test matrix extraction
        matrix = get_matrix(tensor_field, 1, 1)  # At grid point (1,1)
        @test matrix ≈ [1.0 2.0; 3.0 4.0]
        
        # Test matrix setting
        new_matrix = [5.0 6.0; 7.0 8.0]
        set_matrix!(tensor_field, new_matrix, 1, 1)
        extracted = get_matrix(tensor_field, 1, 1)
        @test extracted ≈ new_matrix
    end
end

@testset "Tensor Operations" begin
    grid_2d = UnitGrid([11, 11])
    
    @testset "Trace" begin
        tensor_field = Tensor2Field(pos -> [2.0 1.0; 1.0 3.0], grid_2d)
        trace_field = trace(tensor_field)
        
        @test trace_field isa ScalarField
        @test all(data(trace_field) .≈ 5.0)  # 2 + 3 = 5
    end
    
    @testset "Determinant" begin
        tensor_field = Tensor2Field(pos -> [2.0 1.0; 1.0 3.0], grid_2d)
        det_field = det(tensor_field)
        
        @test det_field isa ScalarField
        @test all(data(det_field) .≈ 5.0)  # 2*3 - 1*1 = 5
    end
    
    @testset "Transpose" begin
        tensor_field = Tensor2Field(pos -> [1.0 2.0; 3.0 4.0], grid_2d)
        transposed = transpose(tensor_field)
        
        @test all(data(transposed)[:, :, 1, 1] .≈ 1.0)  # Same
        @test all(data(transposed)[:, :, 1, 2] .≈ 3.0)  # Swapped
        @test all(data(transposed)[:, :, 2, 1] .≈ 2.0)  # Swapped
        @test all(data(transposed)[:, :, 2, 2] .≈ 4.0)  # Same
    end
    
    @testset "Symmetry" begin
        symmetric_tensor = Tensor2Field(pos -> [1.0 2.0; 2.0 3.0], grid_2d)
        asymmetric_tensor = Tensor2Field(pos -> [1.0 2.0; 3.0 4.0], grid_2d)
        
        @test is_symmetric(symmetric_tensor)
        @test !is_symmetric(asymmetric_tensor)
        
        # Test symmetrization
        symmetrized = symmetrize(asymmetric_tensor)
        @test is_symmetric(symmetrized)
    end
    
    @testset "Eigenvalues and Eigenvectors" begin
        # Create diagonal tensor for predictable eigenvalues
        tensor_field = Tensor2Field(pos -> [3.0 0.0; 0.0 1.0], grid_2d)
        
        eigenvals_fields = eigenvals(tensor_field)
        @test length(eigenvals_fields) == 2
        @test all(eigenvals_fields) do field
            field isa ScalarField
        end
        
        # For diagonal matrix [3,0;0,1], eigenvalues should be 3 and 1
        @test all(data(eigenvals_fields[1]) .≈ 3.0)  # First eigenvalue (sorted desc)
        @test all(data(eigenvals_fields[2]) .≈ 1.0)  # Second eigenvalue
        
        # Test eigen (both values and vectors)
        eigvals_fields, eigvecs_fields = eigen(tensor_field)
        @test length(eigvecs_fields) == 2
        @test all(eigvecs_fields) do field
            field isa VectorField
        end
    end
end

@testset "Tensor-Vector Products" begin
    grid_2d = UnitGrid([11, 11])
    
    tensor_field = Tensor2Field(pos -> [2.0 1.0; 0.0 3.0], grid_2d)
    vector_field = VectorField(pos -> [1.0, 2.0], grid_2d)
    
    @testset "Tensor-Vector Product" begin
        result = dot(tensor_field, vector_field)
        @test result isa VectorField
        
        # [2 1; 0 3] * [1; 2] = [4; 6]
        @test all(data(result)[:, :, 1] .≈ 4.0)
        @test all(data(result)[:, :, 2] .≈ 6.0)
    end
    
    @testset "Vector-Tensor Product" begin
        result = dot(vector_field, tensor_field)
        @test result isa VectorField
        
        # [1 2] * [2 1; 0 3] = [2 7]
        @test all(data(result)[:, :, 1] .≈ 2.0)
        @test all(data(result)[:, :, 2] .≈ 7.0)
    end
    
    @testset "Double Contraction" begin
        tensor1 = Tensor2Field(pos -> [1.0 2.0; 3.0 4.0], grid_2d)
        tensor2 = Tensor2Field(pos -> [2.0 1.0; 1.0 2.0], grid_2d)
        
        result = double_dot(tensor1, tensor2)
        @test result isa ScalarField
        
        # Sum of element-wise products: 1*2 + 2*1 + 3*1 + 4*2 = 15
        @test all(data(result) .≈ 15.0)
    end
    
    @testset "Outer Product" begin
        v1 = VectorField(pos -> [1.0, 2.0], grid_2d)
        v2 = VectorField(pos -> [3.0, 4.0], grid_2d)
        
        result = outer_product(v1, v2)
        @test result isa Tensor2Field
        
        # [1; 2] ⊗ [3 4] = [3 4; 6 8]
        @test all(data(result)[:, :, 1, 1] .≈ 3.0)
        @test all(data(result)[:, :, 1, 2] .≈ 4.0)
        @test all(data(result)[:, :, 2, 1] .≈ 6.0)
        @test all(data(result)[:, :, 2, 2] .≈ 8.0)
    end
end

@testset "FieldCollection Tests" begin
    grid_2d = UnitGrid([21, 21])
    
    @testset "Construction" begin
        # Test basic construction
        scalar1 = ScalarField(pos -> sin(pos[1]), grid_2d)
        scalar2 = ScalarField(pos -> cos(pos[2]), grid_2d)
        vector1 = VectorField(pos -> [pos[1], pos[2]], grid_2d)
        
        collection = FieldCollection([scalar1, scalar2, vector1], ["u", "v", "vel"])
        
        @test length(collection) == 3
        @test labels(collection) == ["u", "v", "vel"]
        @test grid(collection) === grid_2d
        
        # Test access by index
        @test collection[1] === scalar1
        @test collection[2] === scalar2
        @test collection[3] === vector1
        
        # Test access by label
        @test collection["u"] === scalar1
        @test collection["v"] === scalar2
        @test collection["vel"] === vector1
        
        @test_throws KeyError collection["nonexistent"]
    end
    
    @testset "Specialized Constructors" begin
        # Test reaction-diffusion system
        rd_fields = reaction_diffusion_fields(grid_2d, 3)
        @test length(rd_fields) == 3
        @test all(field isa ScalarField for field in rd_fields.fields)
        @test labels(rd_fields) == ["species_1", "species_2", "species_3"]
        
        # Test fluid dynamics
        fluid_fields = fluid_dynamics_fields(grid_2d)
        @test length(fluid_fields) == 2
        @test fluid_fields["density"] isa ScalarField
        @test fluid_fields["velocity"] isa VectorField
        
        # Test elasticity
        elastic_fields = elasticity_fields(grid_2d)
        @test length(elastic_fields) == 2
        @test elastic_fields["displacement"] isa VectorField
        @test elastic_fields["stress"] isa Tensor2Field
    end
    
    @testset "Field Type Filtering" begin
        scalar1 = ScalarField(Float64, grid_2d)
        scalar2 = ScalarField(Float64, grid_2d)
        vector1 = VectorField(Float64, grid_2d)
        tensor1 = Tensor2Field(Float64, grid_2d)
        
        mixed_collection = FieldCollection([scalar1, vector1, scalar2, tensor1])
        
        scalar_fields, scalar_indices = get_scalar_fields(mixed_collection)
        @test length(scalar_fields) == 2
        @test scalar_indices == [1, 3]
        
        vector_fields, vector_indices = get_vector_fields(mixed_collection)
        @test length(vector_fields) == 1
        @test vector_indices == [2]
        
        tensor_fields, tensor_indices = get_tensor_fields(mixed_collection)
        @test length(tensor_fields) == 1
        @test tensor_indices == [4]
    end
    
    @testset "Mathematical Operations" begin
        # Create two compatible collections
        fields1 = [ScalarField(pos -> 1.0, grid_2d), VectorField(pos -> [2.0, 3.0], grid_2d)]
        fields2 = [ScalarField(pos -> 2.0, grid_2d), VectorField(pos -> [1.0, 1.0], grid_2d)]
        
        collection1 = FieldCollection(fields1, ["s", "v"])
        collection2 = FieldCollection(fields2, ["s", "v"])
        
        # Test addition
        sum_collection = collection1 + collection2
        @test all(data(sum_collection[1]) .≈ 3.0)  # 1 + 2
        @test all(data(sum_collection[2])[:, :, 1] .≈ 3.0)  # 2 + 1
        @test all(data(sum_collection[2])[:, :, 2] .≈ 4.0)  # 3 + 1
        
        # Test scalar multiplication
        scaled_collection = collection1 * 2.0
        @test all(data(scaled_collection[1]) .≈ 2.0)
        @test all(data(scaled_collection[2])[:, :, 1] .≈ 4.0)
        @test all(data(scaled_collection[2])[:, :, 2] .≈ 6.0)
    end
    
    @testset "Analysis Functions" begin
        scalar_field = ScalarField(pos -> pos[1]^2 + pos[2]^2, grid_2d)
        vector_field = VectorField(pos -> [pos[1], pos[2]], grid_2d)
        
        collection = FieldCollection([scalar_field, vector_field], ["scalar", "vector"])
        
        # Test statistics
        stats = get_statistics(collection)
        @test haskey(stats, :fields)
        @test haskey(stats, :collection_norm)
        @test haskey(stats, :n_fields)
        @test stats.n_fields == 2
        
        # Test extrema
        extrema_list = get_extrema(collection)
        @test length(extrema_list) == 2
        
        # Test integration
        integrals = integrate(collection)
        @test length(integrals) == 2
    end
end

@testset "Field Conversion and Analysis" begin
    grid_2d = UnitGrid([21, 21])
    
    @testset "Vector to Scalar Conversion" begin
        vector_field = VectorField(pos -> [3.0, 4.0], grid_2d)
        
        # Test magnitude conversion
        mag_field = to_scalar(vector_field, :magnitude)
        @test all(data(mag_field) .≈ 5.0)
        
        # Test component extraction
        x_comp = to_scalar(vector_field, :x)
        @test all(data(x_comp) .≈ 3.0)
        
        y_comp = to_scalar(vector_field, :y)
        @test all(data(y_comp) .≈ 4.0)
        
        # Test max/min operations
        max_field = to_scalar(vector_field, :max)
        @test all(data(max_field) .≈ 4.0)
        
        min_field = to_scalar(vector_field, :min)
        @test all(data(min_field) .≈ 3.0)
    end
    
    @testset "Tensor to Scalar Conversion" begin
        tensor_field = Tensor2Field(pos -> [2.0 1.0; 1.0 3.0], grid_2d)
        
        # Test Frobenius norm
        frob_field = to_scalar(tensor_field, :frobenius_norm)
        expected_norm = sqrt(2^2 + 1^2 + 1^2 + 3^2)  # sqrt(15) ≈ 3.873
        @test all(abs.(data(frob_field) .- expected_norm) .< 1e-10)
        
        # Test trace
        trace_field = to_scalar(tensor_field, :trace)
        @test all(data(trace_field) .≈ 5.0)
        
        # Test determinant
        det_field = to_scalar(tensor_field, :determinant)
        @test all(data(det_field) .≈ 5.0)  # 2*3 - 1*1 = 5
    end
end

@testset "Integration with Grid Systems" begin
    # Test different grid types to ensure compatibility
    grids = [
        UnitGrid([21]),
        UnitGrid([21, 21]),
        UnitGrid([11, 11, 11])
    ]
    
    @testset "Grid Compatibility" begin
        for (i, grid) in enumerate(grids)
            N = ndim(grid)
            
            # Test that fields can be created on each grid type
            scalar_field = ScalarField(Float64, grid)
            @test grid(scalar_field) === grid
            
            vector_field = VectorField(Float64, grid)
            @test grid(vector_field) === grid
            @test size(data(vector_field))[end] == N
            
            tensor_field = Tensor2Field(Float64, grid)
            @test grid(tensor_field) === grid
            @test size(data(tensor_field))[end-1:end] == (N, N)
        end
    end
end

@testset "Performance and Memory Tests" begin
    # Test that operations don't create excessive temporary allocations
    grid_2d = UnitGrid([21, 21])
    
    @testset "Memory Efficiency" begin
        v1 = VectorField(pos -> [1.0, 2.0], grid_2d)
        v2 = VectorField(pos -> [3.0, 4.0], grid_2d)
        
        # Test that in-place operations don't allocate unnecessarily
        v_copy = copy(v1)
        original_data = data(v_copy)
        
        v_copy += v2  # In-place operation
        @test data(v_copy) === original_data  # Same memory location
        
        # Test scalar operations
        v_copy *= 2.0
        @test data(v_copy) === original_data  # Still same memory
    end
end

end  # Main test set